import os
import sys
sys.path.append('./')
import argparse
import torch
import joblib
from smplx import SMPLX

import config as _C
from vis.animation import render_scene_list, renderables

import os
os.environ['PYOPENGL_PLATFORM'] = 'osmesa'

if __name__ == "__main__":
    
    parser = argparse.ArgumentParser()
    parser.add_argument('-s', '--sequence', default='')
    args = parser.parse_args()

    if args.sequence != '':
        _C.SEQUENCE_NAME = args.sequence

    annotation_pth = os.path.join(_C.DATA_BASE_DIR, args.sequence, f"{args.sequence}_annot.pkl")
    annotation = joblib.load(annotation_pth)
    gender = annotation["gender"]
    
    # aitviewer params
    to_tensor = lambda x: torch.from_numpy(x).float()
    kwargs = {
        'poses_root': to_tensor(annotation['global_orient']),
        'poses_body': to_tensor(annotation['body_pose']),
        'betas': to_tensor(annotation['betas']),
        'trans': to_tensor(annotation['transl']),
    }

    smplx = SMPLX(_C.SMPLX_MODEL_DIR, 
                  num_betas=len(annotation['betas'][0]), 
                  gender=gender,
    )
    scene_list = []
    scene_list.append(renderables.addSMPLSequence('smplx', z_up=True, color=[0.65, 0.65, 0.65, 0.8], bm=smplx, **kwargs))
    render_scene_list(scene_list)